import torch.nn as nn

class angle_decoder(nn.Module):
    '''
    基于 Hg 输出的 H/4 x W/4 x 256 的特征图
    '''

    def __init__(self):
        super(angle_decoder, self).__init__()
        # 角度预测单独分支
        self.angle_predict = nn.Sequential(
            nn.Conv2d(256, 256, kernel_size=3, padding=1, stride=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.PixelShuffle(2),
            nn.Conv2d(64, 64, kernel_size=3, padding=1, stride=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.PixelShuffle(2),
            nn.Conv2d(16, 1, kernel_size=1, padding=0, stride=1),
            nn.ReLU6(inplace=True),
        )

    # 角度预测单独分支
    def forward(self, x):
        y = self.angle_predict(x)
        return y